import unittest
import sys
import pymc as pm
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

import thermochron.detrital_ui.run as run
import thermochron.detrital_ui.common as common


class ModelSetup(object):
    def __init__(self):
        #Data
        self.dt_age_file="../data/InyoAges.txt"
        self.hypsometry_file="../data/Inyo.xyhs"
        self.elevation_column='h'
        self.br_file = "../data/bedrock.csv"
        #Model name
        self.model_name="Inyo"
        #Number of breaks
        self.breaks=0
        #MCMC iterations
        self.iterations=50000
        #MCMC final chain size
        self.finalChainSize=50
        #Priors
        self.erate_prior = (0.001, 1)
        self.hc_prior = (0, 2)
        self.abr_prior = (0, 60)
        self.error_prior = (0.11, 0.12) 
        
class TestRun(unittest.TestCase):
    def setUp(self):
        #Avg e1/hc from paper
         #self.avg_e1=0.057
         #self.avg_hc=0.59
        #Stdev from paper
         #self.sd_e1=0.006
         #self.sd_hc=0.22
        
        self.settings = ModelSetup()
        self.M=run.run_MCMC(self.settings)

    
        self.M_avg_e1=np.mean(self.M.trace('e1')[:])
        self.M_avg_hc=np.mean(self.M.trace('hc')[:])
        #self.M_sd_e1=np.std(self.M.trace('e1')[:])
        #self.M_sd_hc=np.std(self.M.trace('hc')[:])
        
        self.chain_e1=np.genfromtxt("Inyo_0brk_benchmark.txt/Chain_0/e1.txt", comments='#')
        self.chain_hc=np.genfromtxt("Inyo_0brk_benchmark.txt/Chain_0/hc.txt", comments='#')

        self.subset_e1=common.sample_wr(self.chain_e1, self.settings.finalChainSize)
        self.subset_hc=common.sample_wr(self.chain_hc, self.settings.finalChainSize)

        
    #Kolmogorov-Smirnov test
    
    def test_ks_e1(self):
        D, P = sp.stats.ks_2samp(self.M.trace('e1')[:], self.subset_e1)
        print "\nP-value (e1) = ", P
        self.assertTrue(P>0.1)

        D, P = sp.stats.ks_2samp(self.M.trace('hc')[:], self.subset_hc)
        print "\nP-value (hc) = ", P
        self.assertTrue(P>0.1)

        self.assertAlmostEqual(self.M_avg_e1, np.mean(self.subset_e1), places=2)

        self.assertAlmostEqual(self.M_avg_hc, np.mean(self.subset_hc), places=1)
        
##    def test_ks_hc(self):
##        D, P = ss.stats.ks_2samp(self.M.trace('hc')[:], self.subset_hc)
##        print P
##        self.assertTrue(P>0.1)
##        
##    def test_Avg_e1(self):
##        self.assertAlmostEqual(self.M_avg_e1, np.mean(self.subset_e1), places=2)
##
##    def test_Avg_hc(self):
##        self.assertAlmostEqual(self.M_avg_hc, np.mean(self.subset_hc), places=1)
        
## Testing against values in the paper   
##    def test_Avg_e1(self):
##        self.assertAlmostEqual(self.avg_e1, self.M_avg_e1, places=2)
##        self.assertAlmostEqual(self.sd_e1, self.M_sd_e1, places=2)
##        self.assertAlmostEqual(self.avg_hc, self.M_avg_hc, places=1)
##        self.assertAlmostEqual(self.sd_hc, self.M_sd_hc, places=1)
##        
##    def test_StDev_e1(self):
##        self.assertAlmostEqual(self.sd_e1, self.M_sd_e1, places=2)
##
##    def test_Avg_hc(self):
##        self.assertAlmostEqual(self.avg_hc, self.M_avg_hc, places=1)
##        
##    def test_StDev_hc(self):
##        self.assertAlmostEqual(self.sd_hc, self.M_sd_hc, places=1)

    
        
if __name__ == '__main__':
    unittest.main()

                             
